from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request, HTTPConnection
from starlette.responses import Response
from starlette.types import ASGIApp, Receive, Send, Scope

from common.pkg.str_id import get_short_uuid, get_id_with_date_info


class TraceMiddleware(BaseHTTPMiddleware):
    def __init__(self, app: ASGIApp):
        super().__init__(app)
        self.request_id = None

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        response = await call_next(request)
        response.headers.setdefault("REQUEST_ID", self.request_id)
        return response

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] not in ["http", "websocket"]:
            await self.app(scope, receive, send)
            return

        conn = HTTPConnection(scope)
        request_id = conn.headers.get("REQUEST_ID")
        if not request_id:
            request_id = f"{get_id_with_date_info(prefix='trace_')}_{get_short_uuid()}"
        conn.state.request_id = self.request_id = request_id
        await self.app(scope, receive, send)



